﻿using System;
using System.Linq;
using System.Runtime.InteropServices;
using Ast.Core.Services.Extensions;
using Ast.Core.WinApi;
using Ast.Framework;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface IKernelRoutineService : IDependency
    {
        void ListProcessNotifyRoutines(Func<KernelRoutineService.NotifyRoutine, bool> listProcFunc);
        void ListThreadNotifyRoutines(Func<KernelRoutineService.NotifyRoutine, bool> listProcFunc);
        void ListImageNotifyRoutines(Func<KernelRoutineService.NotifyRoutine, bool> listProcFunc);
        void ListRegisterCallbacks(Func<KernelRoutineService.CmRegisterCallbackListItem, bool> listProcFunc);
        void ListShutdownNotifications(Func<KernelRoutineService.ShutdownNotification, bool> listProcFunc);
        void ListBugCheckCallbacks(Func<KernelRoutineService.KBugCheckCallbackRecord, bool> listProcFunc);
        void ListBugCheckReasonCallbacks(Func<KernelRoutineService.KBugCheckReasonCallbackRecord, bool> listProcFunc);
        void ListFsRegistrationChangeRoutines(Func<KernelRoutineService.FsRegistrationChangeRoutine, bool> listProcFunc);
        void ListPowerSettingCallbacks(Func<KernelRoutineService.PowerSettingCallback, bool> listProcFunc);
        void ListPlugPlayNotifications(Func<KernelRoutineService.PlugPlayNotification, bool> listProcFunc);
        void ListLogonSessionTerminatedRoutines(Func<KernelRoutineService.LogonSessionTerminatedRoutine, bool> listProcFunc);
        IntPtr GetLegoNotifyRoutine();
    }

    public class KernelRoutineService : IKernelRoutineService
    {
        private readonly ISymbolAddressFinder _symbolAddressFinder;
        private readonly IDeviceController _deviceController;

        public KernelRoutineService(ISymbolAddressFinder symbolAddressFinder, IDeviceController deviceController)
        {
            _symbolAddressFinder = symbolAddressFinder;
            _deviceController = deviceController;

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public void ListProcessNotifyRoutines(Func<NotifyRoutine, bool> listProcFunc)
        {
            ListNotifies("PspCreateProcessNotifyRoutine", listProcFunc, 64);
        }

        public void ListThreadNotifyRoutines(Func<NotifyRoutine, bool> listProcFunc)
        {
            ListNotifies("PspCreateThreadNotifyRoutine", listProcFunc, 64);
        }

        public void ListImageNotifyRoutines(Func<NotifyRoutine, bool> listProcFunc)
        {
            ListNotifies("PspLoadImageNotifyRoutine", listProcFunc, 8);
        }

        private void ListNotifies(string symbolName, Func<NotifyRoutine, bool> listProcFunc, int size)
        {
            var listAdr = _symbolAddressFinder.FindAddressInNtModule(symbolName);
            var itemAdrs = _deviceController.ReadArrayAs<IntPtr>(listAdr, size)
                .Where(x => x != IntPtr.Zero)
                .Select(x => new IntPtr(x.ToInt32() & ~7));
            foreach (var itemAdr in itemAdrs)
            {
                var notifyRoutine = _deviceController.ReadAs<NotifyRoutine>(itemAdr);
                var goOn = listProcFunc(notifyRoutine);
                if (!goOn) break;
            }
        }

        public void ListRegisterCallbacks(Func<CmRegisterCallbackListItem, bool> listProcFunc)
        {
            ListCallbacks("CallbackListHead", listProcFunc);
        }

        public void ListShutdownNotifications(Func<ShutdownNotification, bool> listProcFunc)
        {
            ListCallbacks("IopNotifyShutdownQueueHead", listProcFunc);
        }

        public void ListBugCheckCallbacks(Func<KBugCheckCallbackRecord, bool> listProcFunc)
        {
            ListCallbacks("KeBugCheckCallbackListHead", listProcFunc);
        }

        public void ListBugCheckReasonCallbacks(Func<KBugCheckReasonCallbackRecord, bool> listProcFunc)
        {
            ListCallbacks("KeBugCheckReasonCallbackListHead", listProcFunc);
        }

        public void ListFsRegistrationChangeRoutines(Func<FsRegistrationChangeRoutine, bool> listProcFunc)
        {
            ListCallbacks("IopFsNotifyChangeQueueHead", listProcFunc);
        }

        public void ListPowerSettingCallbacks(Func<PowerSettingCallback, bool> listProcFunc)
        {
            ListCallbacks("PopRegisteredPowerSettingCallbacks", listProcFunc);
        }

        private void ListCallbacks<T>(string symbolName, Func<T, bool> listProcFunc) where T : IListable
        {
            var listHeadAdr = _symbolAddressFinder.FindAddressInNtModule(symbolName);
            var listHead = _deviceController.ReadAsIntPtr(listHeadAdr).Value;
            var current = _deviceController.ReadAs<T>(listHead);
            while (current.ListEntry.Flink != listHead)
            {
                var goOn = listProcFunc(current);
                if (!goOn) break;

                // next
                current = _deviceController.ReadAs<T>(current.ListEntry.Flink);
            }
        }

        public void ListPlugPlayNotifications(Func<PlugPlayNotification, bool> listProcFunc)
        {
            var pnpDeviceClassNotifyList = _symbolAddressFinder.FindAddressInNtModule("PnpDeviceClassNotifyList");
            var listEntries = _deviceController.ReadArrayAs<Structure.ListEntry>(pnpDeviceClassNotifyList, 13); // an array having 13 elements
            for (var index = 0; index < listEntries.Length; index++)
            {
                var listHead = pnpDeviceClassNotifyList + index*8;
                var currentListEntry = listEntries[index];

                var goOn = true;

                while (currentListEntry.Flink != listHead)
                {
                    var currentPlugPlayNotification = _deviceController.ReadAs<PlugPlayNotification>(currentListEntry.Flink);

                    goOn = listProcFunc(currentPlugPlayNotification);
                    if (!goOn) break;

                    // next
                    currentListEntry = currentPlugPlayNotification.Entry;
                }

                if (!goOn) break;
            }
        }

        public void ListLogonSessionTerminatedRoutines(Func<LogonSessionTerminatedRoutine, bool> listProcFunc)
        {
            var seFileSystemNotifyRoutinesHeadPtr = _symbolAddressFinder.FindAddressInNtModule("SeFileSystemNotifyRoutinesHead");
            var seFileSystemNotifyRoutinesHead = _deviceController.ReadAsIntPtr(seFileSystemNotifyRoutinesHeadPtr).Value;
            var currentLogonSessionTerminatedRoutine = seFileSystemNotifyRoutinesHead;
            while (currentLogonSessionTerminatedRoutine != IntPtr.Zero)
            {
                var logonSessionTerminatedRoutine = _deviceController.ReadAs<LogonSessionTerminatedRoutine>(currentLogonSessionTerminatedRoutine);

                var goOn = listProcFunc(logonSessionTerminatedRoutine);
                if (!goOn) break;

                // next
                currentLogonSessionTerminatedRoutine = logonSessionTerminatedRoutine.Next;
            }
        }

        public IntPtr GetLegoNotifyRoutine()
        {
            var pspLegoNotifyRoutinePtr = _symbolAddressFinder.FindAddressInNtModule("PspLegoNotifyRoutine");
            return _deviceController.ReadAsIntPtr(pspLegoNotifyRoutinePtr).Value;
        }

        public interface IListable
        {
            Structure.ListEntry ListEntry { get; }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct CmRegisterCallbackListItem : IListable
        {
            [MarshalAs(UnmanagedType.Struct)]
            public Structure.ListEntry List;

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 8)]
            private byte[] _;

            public long Cookie;

            public IntPtr Context;

            public IntPtr Function;

            [MarshalAs(UnmanagedType.Struct)]
            public Structure.UnicodeString Altitude;

            public Structure.ListEntry ListEntry
            {
                get { return List; }
            }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct NotifyRoutine
        {
            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 4)]
            private byte[] _;

            public IntPtr Function;
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct ShutdownNotification : IListable
        {
            public Structure.ListEntry List;

            public IntPtr DeviceObject;

            public Structure.ListEntry ListEntry
            {
                get { return List; }
            }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct KBugCheckCallbackRecord : IListable
        {
            public Structure.ListEntry Entry;

            public IntPtr CallbackRoutine;

            public Structure.ListEntry ListEntry
            {
                get { return Entry; }
            }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct KBugCheckReasonCallbackRecord : IListable
        {
            public Structure.ListEntry Entry;

            public IntPtr CallbackRoutine;

            public Structure.ListEntry ListEntry
            {
                get { return Entry; }
            }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct PlugPlayNotification
        {
            public Structure.ListEntry Entry;

            public int EventCategory;

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x8)]
            private byte[] _;

            public IntPtr CallbackRoutine;

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x4)]
            private byte[] __;

            public IntPtr DriverObject;
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct LogonSessionTerminatedRoutine
        {
            public IntPtr Next;

            public IntPtr CallbackRoutine;
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct FsRegistrationChangeRoutine : IListable
        {
            public Structure.ListEntry Entry;

            public IntPtr DriverObject;

            public IntPtr DriverNotificationRoutine;

            public Structure.ListEntry ListEntry
            {
                get { return Entry; }
            }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct PowerSettingCallback : IListable
        {
            public Structure.ListEntry Entry;

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x20)]
            private byte[] _;

            public IntPtr Callback;

            public Structure.ListEntry ListEntry
            {
                get { return Entry; }
            }
        }
    }
}

/*
typedef struct _KBUGCHECK_REASON_CALLBACK_RECORD {
    LIST_ENTRY Entry;
    PKBUGCHECK_REASON_CALLBACK_ROUTINE CallbackRoutine;
    PUCHAR Component;
    ULONG_PTR Checksum;
    KBUGCHECK_CALLBACK_REASON Reason;
    UCHAR State;
} KBUGCHECK_REASON_CALLBACK_RECORD, *PKBUGCHECK_REASON_CALLBACK_RECORD;

typedef struct _KBUGCHECK_CALLBACK_RECORD {
    LIST_ENTRY Entry;
    PKBUGCHECK_CALLBACK_ROUTINE CallbackRoutine;
    _Field_size_bytes_opt_(Length) PVOID Buffer;
    ULONG Length;
    PUCHAR Component;
    ULONG_PTR Checksum;
    UCHAR State;
} KBUGCHECK_CALLBACK_RECORD, *PKBUGCHECK_CALLBACK_RECORD;
*/